import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from copy import deepcopy


def get_nearest_C_num(xj, u):
    """ 获得距离样本xj最近的均值向量的簇标记 """
    distance = np.zeros(len(u))
    for i in range(len(distance)):
        # 计算样本xj与各均值向量的距离
        distance[i] = np.sum((xj - u[i]) ** 2)
    return np.argmin(distance)


def k_means(dataset, k=3):
    """ k-均值算法 """
    # 随机初始化均值向量索引
    print('\nk = {}'.format(k))
    random_int = np.random.randint(0, len(dataset), size=k)
    while len(np.unique(random_int)) < k:
        random_int = np.random.randint(0, len(dataset), size=k)
    
    u = np.zeros([k, 2])  # 储存均值向量
    for i in range(k):
        u[i] = list(dataset.iloc[random_int[i], :-1])
    u_init = deepcopy(u)
    print('初始中心点：')
    print(u_init)

    count = 0
    while 1:
        u_old = deepcopy(u)  # 记录更新前的均值向量
        # 创建并初始化簇划分，令每个簇为空集
        C = dict()
        for i in range(k):
            C[i] = []
        
        for j in range(len(dataset)):
            # 获得距离样本xj最近的均值向量的簇标记
            lambda_j = get_nearest_C_num(np.array(dataset.iloc[j, :-1]), u)
            # 将样本xj划入相应的簇
            C[lambda_j].append(list(dataset.iloc[j, :-1]))

        # 更新均值向量
        for i in range(k):
            u[i] = np.mean(C[i], axis=0)

        count += 1
        # 判断是否达到迭代终止条件
        if np.linalg.norm(u - u_old) < 0.001:
            break

    print('迭代次数：', count)
    color = ['r', 'g', 'b', 'y', 'c', 'm']
    plt.figure()
    # 绘制初始中心点
    plt.scatter(u_init[:, 0], u_init[:, 1], edgecolors='k', facecolors='none', s=100)
    for i in range(k):
        # 绘制簇内样本和均值向量
        plt.plot(np.array(C[i])[:,0], np.array(C[i])[:,1], color[i] + '.')
        plt.plot(u[i,0], u[i,1], color[i] + '+')
    plt.title('k = {}'.format(k))


if __name__ == '__main__':
    # 读取西瓜数据集4.0
    dataset = pd.read_csv('watermelon_9.csv')
    dataset = dataset.drop(['id'], axis=1)

    k_means(dataset, k=3)
    k_means(dataset, k=4)
    k_means(dataset, k=5)

    plt.show()